# !/usr/bin/env python3
# coding=utf8
"""
将通达信的某个代码的收盘价数据导出到 filename 里,
载入 filename 里的收盘价数据, 计算 timeperiod 周期的均线, 输出,
打开通达信的 K 线界面, 比较计算的均线数据和通达信界面上的均线数据, 结果一致表示计算逻辑一致,
"""
import json
import numpy
import random
# import talib
from typing import List, Dict, Tuple, Union


def gen_values(base: float, count: int) -> List[float]:
    """
    以 base 为基准, 每次在涨跌 10% 的范围内波动, 生成 count 个数据,
    逻辑不严谨, 涨停价没有向下靠拢, 跌停价没有向上靠拢, round 函数不是严格的四舍五入,
    因为是随机生成一些假数据, 所以不用那么严谨, 应当以逻辑清晰为主,
    """
    values: list = [float(base)]

    for _ in range(count - 1):
        dn: float = values[-1] * (1 - 0.1)
        up: float = values[-1] * (1 + 0.1)
        # https://zhuanlan.zhihu.com/p/115431517
        value: float = round(random.uniform(a=dn, b=up), 2)
        values.append(value)

    return values


def SMA_self_v1(real: list, timeperiod: int) -> list:
    """
    不依赖三方库, 自己实现的均线函数, 第一版,
    """
    assert isinstance(timeperiod, int) and 0 < timeperiod
    # 如果是5日均线, 就要把最开始的5个值求平均, 所以最前面的4个值是无效的
    dst: list = [None for _ in range(timeperiod - 1)]

    for i in range(timeperiod - 1, len(real), 1):
        arr: list = real[i - (timeperiod - 1):i + 1]
        avg = sum(arr) / len(arr)
        dst.append(avg)

    return dst


def SMA_self_v2(real: list, timeperiod: int) -> list:
    """
    不依赖三方库, 自己实现的均线函数, 第二版,
    """
    assert isinstance(timeperiod, int) and 0 < timeperiod

    dst: list = []

    sum = 0
    for i in range(0, len(real), 1):
        # 5日线的时候, i=4刚好0~4共5天, 不需要弹出旧值,
        pop_value = 0 if i < timeperiod else real[i - timeperiod]
        sum = sum - pop_value + real[i]
        # 5日线的时候, i=4刚好0~4共5天, 值开始有效,
        avg = None if i < timeperiod - 1 else sum / timeperiod
        dst.append(avg)

    return dst


def SMA_numpy(real: Union[list, numpy.ndarray], timeperiod: int) -> numpy.ndarray:
    """
    用 numpy 实现的均线函数,
    """
    if isinstance(real, list):
        real = numpy.array(real, dtype=float)
    assert isinstance(real, numpy.ndarray)

    weights: numpy.ndarray = numpy.ones(shape=timeperiod) / timeperiod
    results: numpy.ndarray = numpy.convolve(a=weights, v=real)
    # numpy中的convolve的理解
    # https://blog.csdn.net/u011599639/article/details/76254442

    dst: numpy.ndarray = results[:-timeperiod + 1]
    for i in range(0, timeperiod - 1):
        dst[i] = None

    return dst


def SMA_talib(real: Union[list, numpy.ndarray], timeperiod: int) -> numpy.ndarray:
    """
    调用 talib 的均线函数,
    Simple Moving Average (Overlap Studies)
    注意: 需求(2 <= timeperiod)因为(1 == timeperiod)时会报错,
    """
    if isinstance(real, list):
        real = numpy.array(real, dtype=float)
    assert isinstance(real, numpy.ndarray)

    import talib
    return talib.SMA(real, timeperiod)


def SMA_test(real: list, timeperiod: int):
    """
    测试函数, 检查各均线函数的结果是否一致, 只要函数没抛异常, 测试就通过了,
    """
    def fmt(value) -> str:
        """"""
        if value is None:
            return str(numpy.nan)
        elif isinstance(value, (int, float)):
            return "{:.6f}".format(value)
        elif isinstance(value, numpy.float64):
            if numpy.isnan(value):
                return str(value)
            else:
                return "{:.6f}".format(value)
        else:
            raise ValueError(f"illegal, value_type={type(value)}, value_data={value},")

    result_1 = SMA_self_v1(real=real, timeperiod=timeperiod)
    result_2 = SMA_self_v2(real=real, timeperiod=timeperiod)
    result_3 = SMA_numpy(real=real, timeperiod=timeperiod)
    result_4 = SMA_talib(real=real, timeperiod=timeperiod)

    assert len(result_1) == len(result_2) == len(result_3) == len(result_4)

    for i in range(0, len(result_1), 1):
        item_1 = result_1[i]
        item_2 = result_2[i]
        item_3 = result_3[i]
        item_4 = result_4[i]
        assert fmt(item_1) == fmt(item_2) == fmt(item_3) == fmt(item_4)

    return True


def calculate_tdx_MA(filename: str, timeperiod: int):
    """
    将通达信的某个代码的收盘价数据导出到 filename 里,
    载入 filename 里的收盘价数据, 计算 timeperiod 周期的均线, 输出,
    打开通达信的 K 线界面, 比较计算的均线数据和通达信界面上的均线数据, 结果一致表示计算逻辑一致,

    经验证, "通达信"的函数 MA 就是 talib.SMA
    其中, 对于"通达信"的公式"MA1:MA(CLOSE,M1);"有如下提示:
    返回简单移动平均
    用法:
    MA(X,N):X的N日简单移动平均,算法(X1+X2+X3+...+Xn)/N,N支持变量
    """
    with open(file=filename, mode="r", encoding="utf8") as f:
        dd_dict: Dict[str, float] = json.load(fp=f)

    ll_list: List[Tuple[str, float]] = [(dt, px) for dt, px in dd_dict.items()]
    ll_list.sort(key=lambda elem: elem[0], reverse=False)

    dt_list: list = [elem[0] for elem in ll_list]  # 日期列表
    px_list: list = [elem[1] for elem in ll_list]  # 价格列表

    SMA_test(real=px_list, timeperiod=timeperiod)  # 检查各均线函数的结果是否一致

    sma_arr: numpy.ndarray = SMA_talib(real=px_list, timeperiod=timeperiod)

    assert len(dt_list) == len(sma_arr)

    sma_map: Dict[str, numpy.float64] = dict(zip(dt_list, sma_arr))

    for dt, px in sma_map.items():
        print(dt, "{:.2f}".format(px))


if __name__ == '__main__':
    import os
    filename: str = r'tdx_SSE_000300.json'
    filepath: str = os.path.join(os.path.dirname(os.path.abspath(__file__)), filename)
    timeperiod: int = 2
    calculate_tdx_MA(filepath, timeperiod)
